import os
import onnx
import csv
from onnxruntime.quantization import quantize_static, QuantFormat, QuantType, CalibrationMethod, CalibrationDataReader

# ========== 1. 数据预处理与DataReader ==========
def preprocess_image(image_path, target_size=(640, 640)):
    # TODO: 按需实现，返回模型输入格式的numpy数组
    pass

class DataReader(CalibrationDataReader):
    def __init__(self, image_paths, batch_size=1):
        self.datas = [preprocess_image(p) for p in image_paths]
        self.batch_size = batch_size
        self.idx = 0

    def get_next(self):
        if self.idx >= len(self.datas):
            return None
        batch = self.datas[self.idx:self.idx+self.batch_size]
        self.idx += self.batch_size
        # 假设模型输入名为'images'
        return {'images': batch[0] if self.batch_size == 1 else np.concatenate(batch, 0)}

# ========== 2. 节点/类型遍历工具 ==========
def get_all_nodes(model_path):
    model = onnx.load(model_path)
    return [(node.name, node.op_type) for node in model.graph.node]

def group_nodes_by_type(nodes):
    from collections import defaultdict
    op_type_to_nodes = defaultdict(list)
    for name, op_type in nodes:
        op_type_to_nodes[op_type].append(name)
    return op_type_to_nodes

# ========== 3. 量化与评估 ==========
def quantize_and_evaluate(
    model_fp32, nodes_to_exclude, data_reader, quant_model_path, 
    eval_func, eval_args, quant_params
):
    quantize_static(
        model_input=model_fp32,
        model_output=quant_model_path,
        calibration_data_reader=data_reader,
        **quant_params,
        nodes_to_exclude=nodes_to_exclude,
    )
    # 加载量化模型并评估
    return eval_func(quant_model_path, *eval_args)

# ========== 4. 主流程 ==========
def sensitive_analysis(
    model_fp32, image_dir, eval_func, eval_args, 
    mode='type',  # 'type' or 'node'
    csv_path='sensitive_analysis.csv',
    quant_params=None,
    structure_exclude=None,
    batch_size=1,
    max_images=100
):
    nodes = get_all_nodes(model_fp32)
    op_type_to_nodes = group_nodes_by_type(nodes)
    image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir)][:max_images]
    data_reader = DataReader(image_paths, batch_size)
    quant_params = quant_params or dict(
        quant_format=QuantFormat.QDQ,
        activation_type=QuantType.QInt8,
        weight_type=QuantType.QInt8,
        calibrate_method=CalibrationMethod.MinMax,
        per_channel=False,
        reduce_range=False,
    )
    structure_exclude = structure_exclude or []

    # 写表头
    with open(csv_path, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=['exclude', 'num_nodes', 'score'])
        writer.writeheader()

    if mode == 'type':
        for op_type, node_names in op_type_to_nodes.items():
            exclude = node_names + structure_exclude
            quant_model_path = f"quant_exclude_{op_type}.onnx"
            score = quantize_and_evaluate(
                model_fp32, exclude, data_reader, quant_model_path, 
                eval_func, eval_args, quant_params
            )
            with open(csv_path, 'a', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=['exclude', 'num_nodes', 'score'])
                writer.writerow({'exclude': op_type, 'num_nodes': len(node_names), 'score': score})
    elif mode == 'node':
        for name, op_type in nodes:
            exclude = [name] + structure_exclude
            quant_model_path = f"quant_exclude_{name}.onnx"
            score = quantize_and_evaluate(
                model_fp32, exclude, data_reader, quant_model_path, 
                eval_func, eval_args, quant_params
            )
            with open(csv_path, 'a', newline='') as f:
                writer = csv.DictWriter(f, fieldnames=['exclude', 'num_nodes', 'score'])
                writer.writerow({'exclude': name, 'num_nodes': 1, 'score': score})

    print(f"敏感层分析结果已保存: {csv_path}")

# ========== 5. 用法示例 ==========
if __name__ == "__main__":
    # 1. 指定模型、数据、评估函数
    model_fp32 = 'your_model.onnx'
    image_dir = 'your_calib_images'
    def eval_func(quant_model_path, *args):
        # TODO: 实现你的评估逻辑，返回分数（如AP50、Top1等）
        return 0.0
    eval_args = ('test.txt', 'label_dir', 'save_vis_dir')  # 按需修改

    # 2. 运行敏感层分析
    sensitive_analysis(
        model_fp32, image_dir, eval_func, eval_args,
        mode='type',  # 或 'node'
        csv_path='sensitive_analysis.csv',
        quant_params=None,  # 可自定义
        structure_exclude=[],  # 可自定义
        batch_size=1,
        max_images=100
    )